# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： bbox_correct
Description :
Author : 'li'
date： 2021/6/21
-------------------------------------------------
Change Activity:
2021/6/21:
-------------------------------------------------
"""
import cv2
import numpy as np

from ml.cv.algorithms.bbox_generate.corner_to_bbox.find_corner_point import VectorAngleFinder
from ml.cv.algorithms.bbox_generate.corner_to_bbox.origin_point_generate import OriginPointGenerator
from ml.cv.algorithms.bbox_generate.corner_to_bbox.util import get_line_cross_point
from ml.cv.algorithms.sort_points.sort_points import resort_points


class BboxCorrector:
    def __init__(self, mask):
        self.mask = mask
        self.contour = BboxCorrector._init_contour(mask)
        self.cw_min, self.cw_max, self.hw_min, self.hw_max = None, None, None, None
        self.clip_text_area_mask = None
        self.bg = None
        self.bg_bbox = None
        self.border_size = None
        self.bg_contour = None
        self.origin_point_generator = None

    def _clip_text_area_mask(self):
        self.clip_text_area_mask = self.mask[self.hw_min:self.hw_max + 1, self.cw_min: self.cw_max + 1]

    def _gen_border_text_area_mask(self):
        h, w = self.clip_text_area_mask.shape[:2]
        border_size = min(h, w) // 2
        if border_size < 10:
            border_size = 10
        self.border_size = border_size
        one_side_size = border_size
        bg = np.zeros((h + 2 * border_size, w + 2 * border_size))
        bg[one_side_size:h + one_side_size, one_side_size:w + one_side_size] = self.clip_text_area_mask
        self.bg = bg
        pass

    def _get_bg_bbox(self):
        bg = self.bg.astype(np.uint8)
        ret, thresh = cv2.threshold(bg, 127, 255, 0)
        contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 0:
            return None
        if len(contours) > 1:
            contours = [self._get_max_size_contours(contours)]
        self.bg_contour = contours[0]
        rect = cv2.minAreaRect(contours[0])
        box = cv2.boxPoints(rect)
        box = resort_points(box)
        self.bg_bbox = box
        self.origin_point_generator = OriginPointGenerator(self.bg_contour[:, 0, :], self.bg_bbox)

    def get_box(self):
        if self.contour is None or self.contour.size < 3:
            return None
        self.cw_min, self.cw_max, self.hw_min, self.hw_max = self.contour[:, 0, 0].min(), \
                                                             self.contour[:, 0, 0].max(), \
                                                             self.contour[:, 0, 1].min(), \
                                                             self.contour[:, 0, 1].max()
        self._clip_text_area_mask()
        self._gen_border_text_area_mask()
        self._get_bg_bbox()
        points = self.get_format_angle_points()
        if points is None:
            return None
        points = points - self.border_size
        points[:, 0] = points[:, 0] + self.cw_min
        points[:, 1] = points[:, 1] + self.hw_min
        return points

    @staticmethod
    def _get_max_size_contours(contours):
        """
        get max siz contours
        """
        return_contours = contours[0]
        length = len(contours)
        for i in range(1, length):
            ret_size = return_contours.shape[0]
            new_size = contours[i].shape[0]
            if ret_size < new_size:
                return_contours = contours[i]
        return return_contours

    @staticmethod
    def _init_contour(mask):
        _, thresh = cv2.threshold(mask, 127, 255, 0)
        contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 0:
            return None
        if len(contours) > 1:
            contours = [BboxCorrector._get_max_size_contours(contours)]
        contour = contours[0]
        return contour

    @staticmethod
    def _enlarge_bbox(enlarge_pixel, bbox):
        """
        enlarge bbox
        :param enlarge_pixel:
        :param bbox:
        :return:
        """

        p1, p2, p3, p4 = bbox[0], bbox[1], bbox[2], bbox[3]
        p1 = p1[0] - enlarge_pixel, p1[1] - enlarge_pixel
        p2 = p2[0] + enlarge_pixel, p2[1] - enlarge_pixel
        p3 = p3[0] + enlarge_pixel, p3[1] + enlarge_pixel
        p4 = p4[0] - enlarge_pixel, p4[1] + enlarge_pixel
        new_bbox = np.array([p1, p2, p3, p4])
        return new_bbox

    @staticmethod
    def _get_line_point(finder, p0):
        l1 = finder.get_max_angle_point() + p0
        l2 = finder.get_min_angle_point() + p0
        return [l1, l2]

    def _get_line1(self):
        p0 = self.origin_point_generator.top_p
        vectors = self.bg_contour[:, 0, :] - p0
        finder = VectorAngleFinder(vectors)
        return self._get_line_point(finder, p0)

    def _get_line2(self):
        p0 = self.origin_point_generator.right_p
        vectors = self.bg_contour[:, 0, :] - p0
        finder = VectorAngleFinder(vectors)
        return self._get_line_point(finder, p0)

    def _get_line3(self):
        p0 = self.origin_point_generator.down_p
        vectors = self.bg_contour[:, 0, :] - p0
        finder = VectorAngleFinder(vectors)
        return self._get_line_point(finder, p0)

    def _get_line4(self):
        p0 = self.origin_point_generator.left_p
        vectors = self.bg_contour[:, 0, :] - p0
        finder = VectorAngleFinder(vectors)
        vectors_angle = finder.vectors_angle
        location = np.where((vectors_angle > 0) & (vectors_angle < 90))  # 左下角
        filter_angles = vectors_angle[location]
        if filter_angles.size > 0:
            lo = np.argmax(filter_angles)
            true_location = location[0][lo]
            l1 = finder.vectors[true_location]
        else:
            location = np.where((vectors_angle > 270))
            filter_angles = vectors_angle[location]
            lo = np.argmax(filter_angles)
            true_location = location[0][lo]
            l1 = finder.vectors[true_location]
        location = np.where((vectors_angle > 270))
        filter_angles = vectors_angle[location]
        if filter_angles.size > 0:
            lo = np.argmin(filter_angles)
            true_location = location[0][lo]
            l2 = finder.vectors[true_location]
        else:
            location = np.where((vectors_angle > 0) & (vectors_angle < 90))
            filter_angles = vectors_angle[location]
            lo = np.argmin(filter_angles)
            true_location = location[0][lo]
            l2 = finder.vectors[true_location]
        l1 = l1 + p0
        l2 = l2 + p0
        return [l1, l2]

    def get_format_angle_points(self):
        lines = []
        l1 = self._get_line1()
        lines.append(l1)
        l2 = self._get_line2()
        lines.append(l2)
        l3 = self._get_line3()
        lines.append(l3)
        l4 = self._get_line4()
        lines.append(l4)
        points = self._get_cross_points(lines)
        return points

    @staticmethod
    def _get_cross_points(lines):
        p0 = get_line_cross_point(lines[-1], lines[0])
        p1 = get_line_cross_point(lines[0], lines[1])
        p2 = get_line_cross_point(lines[1], lines[2])
        p3 = get_line_cross_point(lines[2], lines[-1])
        if p0 is None or p1 is None or p2 is None or p3 is None:
            return None
        return np.array([p0, p1, p2, p3])
